﻿// Copyright (c) 2025 liucl
// Inject is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

#include "log.hpp"

#include <WinSock2.h>
#include <ws2tcpip.h>

#include <wx/listbox.h>

#include <thread>

wxDECLARE_EVENT(EVT_EVENT_LOG, wxCommandEvent);
class udp_server {
  static constexpr int TX_BUFFER_SIZE = 1024 * 10;
  SOCKET socket_ = INVALID_SOCKET;
  sockaddr_in addr_ = {};
  wxEvtHandler *handler_ = nullptr;
  int run_ = 1;

  static void init_winsock_() {
    WSADATA wsaData;
    auto rv = ::WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (rv != 0) {
    }
  }

  void cleanup_() {
    if (socket_ != INVALID_SOCKET) {
      ::closesocket(socket_);
    }
    socket_ = INVALID_SOCKET;
    ::WSACleanup();
  }

public:
  udp_server(wxEvtHandler *handler, const std::string &host, uint16_t port)
      : handler_(handler) {
    init_winsock_();

    addr_.sin_family = PF_INET;
    addr_.sin_port = htons(port);
    addr_.sin_addr.s_addr = INADDR_ANY;
    if (InetPtonA(PF_INET, host.c_str(), &addr_.sin_addr.s_addr) != 1) {
      int last_error = ::WSAGetLastError();
      ::WSACleanup();
    }

    socket_ = ::socket(PF_INET, SOCK_DGRAM, 0);
    if (socket_ == INVALID_SOCKET) {
      int last_error = ::WSAGetLastError();
      ::WSACleanup();
    }

    int option_value = TX_BUFFER_SIZE;
    if (::setsockopt(socket_, SOL_SOCKET, SO_RCVBUF,
                     reinterpret_cast<const char *>(&option_value),
                     sizeof(option_value)) < 0) {
      int last_error = ::WSAGetLastError();
      cleanup_();
    }
    DWORD timeout = 2 * 1000; // 每两秒处理一次
    if (::setsockopt(socket_, SOL_SOCKET, SO_RCVTIMEO,
                     reinterpret_cast<const char *>(&timeout),
                     sizeof(timeout)) < 0) {
      int last_error = ::WSAGetLastError();
      cleanup_();
    }

    if (::bind(socket_, (struct sockaddr *)&addr_, sizeof(addr_)) == -1) {
      int last_error = ::WSAGetLastError();
      cleanup_();
    }

    start_recv();
  }

  ~udp_server() {
    run_ = 2;
    while (run_ == 2)
      ;
    cleanup_();
  }

  void start_recv() {
    std::thread([this]() {
      char buffer[1024] = {};
      while (run_ == 1) {

        sockaddr_in from = {};
        socklen_t fromlen = sizeof(from);
        int n_bytes = ::recvfrom(socket_, buffer, sizeof(buffer), 0,
                                 (struct sockaddr *)&from, &fromlen);
        if (n_bytes > 0) {
          buffer[n_bytes] = '\0';
          wxCommandEvent evt(EVT_EVENT_LOG, wxID_ANY);
          evt.SetString(wxString::FromUTF8(buffer));
          handler_->AddPendingEvent(evt);
        }
      }
      run_ = 3;
    }).detach();
  }
  SOCKET fd() const { return socket_; }

  void send(const char *data, size_t n_bytes) {
    socklen_t tolen = sizeof(struct sockaddr);
    if (::sendto(socket_, data, static_cast<int>(n_bytes), 0,
                 (struct sockaddr *)&addr_, tolen) == -1) {
    }
  }
};

LogControl::LogControl(wxWindow *parent, wxWindowID id)
    : wxControl(parent, id) {
  list_box_ = new wxListBox(this, wxID_ANY);
  auto sizer = new wxBoxSizer(wxVERTICAL);
  auto btn = new wxButton(this, wxID_ANY, wxT("back"));
  sizer->Add(btn, 0, wxEXPAND);
  sizer->Add(list_box_, 1, wxEXPAND);
  SetSizer(sizer);

  btn->Bind(wxEVT_BUTTON, [this](wxCommandEvent &evt) {
    wxCommandEvent uninject_evt(EVT_EVENT_UNINJECT, wxID_ANY);
    AddPendingEvent(uninject_evt);
  });

  server_ = new udp_server(this, "0.0.0.0", 8888);
}

LogControl::~LogControl() { delete server_; }

void LogControl::OnLog(wxCommandEvent &evt) {
  wxString str = evt.GetString();
  if (list_box_->GetCount() >= 20) {
    list_box_->Delete(0);
  }
  list_box_->AppendString(str);
}

wxDEFINE_EVENT(EVT_EVENT_UNINJECT, wxCommandEvent);
wxDEFINE_EVENT(EVT_EVENT_LOG, wxCommandEvent);
// clang-format off
wxBEGIN_EVENT_TABLE(LogControl, wxControl) 
    EVT_COMMAND(wxID_ANY, EVT_EVENT_LOG, LogControl::OnLog)
wxEND_EVENT_TABLE()